#%%
# python -m lilab.OpenLabCluster_train.t1_cluster_name_rename *.clippredpkl
import pickle
import re
import argparse

clippredfile = '/DATA/taoxianming/rat/data/Mix_analysis/SexAgeDay55andzzcWTinAUT_MMFF/result32/semiseq2seq_iter0/output/semisupervise-decSeq-iter2-epoch5/olc-2024-05-23-semiseq2seq-iter2-epoch5_pca12.clippredpkl'

rename_page = """
1	Back to back leaving
2	Facial investigation
3	Head to head approaching
4	Circly sniffing
5	Both rearing and facing away
6	Cross by leaving
7	Both rearing with facial investigation
8	Back to back facing away
9	Pinning with the partner full rotation
10	Being fast chased
11	Rearing with being back to back faced away
12	Sniffing anogenital
13	Pouncing from behind
14	Pouncing from head
15	Being slowly chased
16	Rearing ahead a rearing partner (investigating back)
17	Rearing and being faced away or leaved
18	Watching or approaching the partner from behind
19	Pouncing from behind and partner evading
20	Pinning with the partner partial rotation
21	Rearing with being watched
22	Rearing with the partner in social range
23	Being pinned and full rotation
24	Fast chasing
25	Back to back facing away the rearing partner
26	Being sniffed at anogenital
27	Being pounced from behind
28	Being pounced from head
29	Slow chasing
30	Rearing behind a rearing partner (investigating back)
31	Facing away or leaving a rearing partner
32	Being watched or approached from behind
33	Evading and being pounced from behind
34	Being pinned with partial rotation
35	Watching the rearing partner
36	Sniffing the rearing partner
"""

mutual_page = """
Back to back leave
Facial investigation
Head to head approach
Circly sniff
Both rear and face away
Cross by leave
Both rear with facial investigation
Back to back face away
Pin full rotation
Fast chase
Rear with being back to back faced away
Sniff anogenital
Pounce from behind
Pounce from head
Slow chase
Rear behind a rearing partner with back investigation
Rear when being faced away or leaved
Watch or approach from behind
Pounce from behind and partner evade
Pin partial rotation
Rear with being watched
Rear within social range
"""

cluster_names_re = [s for s in rename_page.split("\n") if len(s)]
pattern_to_rename = re.compile(r'^(\d+)\s+(.*)')
pattern_extract = [pattern_to_rename.search(s).groups() for s in cluster_names_re]

mutual_names_re = [s for s in mutual_page.split("\n") if len(s)]


def main(clippredfile):
    clippreddata = pickle.load(open(clippredfile, 'rb'))
    assert clippreddata['nK_mutual'] + clippreddata['nK_mirror_half'] == len(mutual_names_re)

    for i, name in pattern_extract:
        clippreddata['cluster_names'][int(i)-1] = name
    clippreddata['cluster_names_mutualmerge'] = mutual_names_re
    pickle.dump(clippreddata, open(clippredfile, 'wb'))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('clippredfile', type=str)
    args = parser.parse_args()
    main(args.clippredfile)
